package com.jboltai.capability.vdb.collection;


import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.vdb.JBoltVDB;
import com.jboltai.capability.vdb.collection.field.VDBField;
import com.jboltai.capability.vdb.collection.field.VDBFieldType;
import com.jboltai.capability.vdb.collection.field.VectorIndexType;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 向量数据库集合，自带一个类型为string，名称为id的主键字段，以及 一个类型为string，名称为doc的用于存储文档的字段
 */
public class VDBCollection {

    /**
     * 默认分片数
     */
    public static final Integer DEFAULT_SHARD_NUM = 1;

    /**
     * 默认副本数
     */
    public static final Integer DEFAULT_REPLICA_NUM = 0;


    /**
     * 集合名称
     */
    private String name;

    /**
     * 指定 Collection 的分片数。分片是把大数据集切成多个子数据集。在搜索时，全部分片是并发执行的，分片数量越多，平均耗时越低，但是过多的分片会带来额外开销而影响性能。<br/>
     * 单分片数据量建议控制在300万以内，例如500万向量，可设置2个分片。<br/>
     * 默认为 1。
     */
    private Integer shardNum = null;

    /**
     * 指定 Collection 的副本数。副本数是指每个主分片有多个相同的备份，用来容灾和负载均衡。<br/>
     * 不是所有向量数据库都支持该参数。
     * 默认值是0
     */
    private Integer replicaNum = null;

    /**
     * 指定 Collection 的描述信息。
     */
    private String description;

    private JBoltVDB vdb;



    private List<VDBField> fields = new ArrayList<VDBField>(){
        {
            add(new VDBField("id", VDBFieldType.STRING));
            add(new VDBField("doc", VDBFieldType.STRING));
        }
    };

    public VDBCollection(JBoltVDB vdb) {
        this.vdb = vdb;
    };


    public String getName() {
        return name;
    }

    public VDBCollection setName(String name) {
        this.name = name;
        return this;
    }

    public Integer getShardNum() {
        return shardNum;
    }

    public VDBCollection setShardNum(Integer shardNum) {
        this.shardNum = shardNum;
        return this;
    }

    public Integer getReplicaNum() {
        return replicaNum;
    }

    public VDBCollection setReplicaNum(Integer replicaNum) {
        this.replicaNum = replicaNum;
        return this;
    }

    public String getDescription() {
        return description;
    }

    public VDBCollection setDescription(String description) {
        this.description = description;
        return this;
    }

    /**
     * 添加普通字段
     * @param name
     * @param VDbFieldType
     * @return
     */
    public VDBCollection addNormalField(String name, VDBFieldType VDbFieldType) {
        return addField(new VDBField(name, VDbFieldType));
    }

    /**
     * 添加向量字段，每个集合只能添加一个
     * @param vectorIndexType
     * @return
     */
    public VDBCollection addVectorField(VectorIndexType vectorIndexType) {
        return addField(new VDBField(vectorIndexType, null));
    }


    /**
     * 添加向量字段，每个集合只能添加一个
     * @param vectorIndexType
     * @param vectorIndexParams 向量索引配置参数
     * @return
     */
    public VDBCollection addVectorField(VectorIndexType vectorIndexType, JSONObject vectorIndexParams) {
        return addField(new VDBField(vectorIndexType, vectorIndexParams));
    }


    /**
     * 添加向量字段，每个集合只能添加一个
     * @param vectorIndexType
     * @param vectorIndexParams 向量索引配置参数
     * @return
     */
    public VDBCollection addVectorField(VectorIndexType vectorIndexType, JSONObject vectorIndexParams,  VDBMetricType VDbMetricType) {
        return addField(new VDBField(vectorIndexType, vectorIndexParams, VDbMetricType));
    }


    public VDBCollection addField(VDBField field) {
        if (field == null) {
            throw new IllegalArgumentException("field不能为null");
        }
        if (StrUtil.isBlank(field.getName())) {
            throw new IllegalArgumentException("filed的name属性不能为空");
        }

        for (VDBField f : fields) {

            if (Objects.equals(f.getName(), field.getName()) && Objects.equals(f.getName(), VDBField.DOC_FIELD_NAME)) {
                //doc字段，允许重复定义，更新它的fullText属性
                f.setFullTextSearch(field.isFullTextSearch());
                return this;
            }

            if (f.isPrimaryKey() && field.isPrimaryKey()) {
                throw new RuntimeException("主键字段只能添加一个");
            }
            if (f.getVdbFieldType() == VDBFieldType.VECTOR && field.getVdbFieldType() == VDBFieldType.VECTOR) {
                throw new RuntimeException("向量字段只能添加一个");
            }
            if (f.getName().equals(field.getName())) {
                throw new RuntimeException("字段名[" + field.getName() + "]重复");
            }
        }

        this.fields.add(field);
        return this;
    }

    public List<VDBField> getFields() {
        return fields;
    }
}
